from flask.sessions import SecureCookieSessionInterface
from werkzeug.datastructures import CallbackDict
from flask import session

class CustomSessionInterface(SecureCookieSessionInterface):
    """自定义Session接口，解决Flask 3.1.1中partitioned参数的兼容性问题"""
    
    def save_session(self, app, session, response):
        """保存session，移除不兼容的partitioned参数"""
        if session is None:
            return
        
        domain = self.get_cookie_domain(app)
        path = self.get_cookie_path(app)
        
        # 如果session为空，删除cookie
        if not session:
            if session.modified:
                response.delete_cookie(
                    app.config['SESSION_COOKIE_NAME'],
                    domain=domain,
                    path=path
                )
            return
        
        # 检查是否需要设置cookie
        if not self.should_set_cookie(app, session):
            return
        
        httponly = self.get_cookie_httponly(app)
        secure = self.get_cookie_secure(app)
        samesite = self.get_cookie_samesite(app)
        expires = self.get_expiration_time(app, session)
        val = self.get_signing_serializer(app).dumps(dict(session))
        
        # 设置cookie，不包含partitioned参数
        response.set_cookie(
            app.config['SESSION_COOKIE_NAME'],
            val,
            expires=expires,
            httponly=httponly,
            domain=domain,
            path=path,
            secure=secure,
            samesite=samesite
        )